﻿using System;
using System.Collections.Generic;
using Mono.Cecil.Metadata;
using Mono.Cecil.PE;

using RVA = System.UInt32;

namespace Mono.Cecil.Cil {

	sealed class CodeReader : ByteBuffer {

		int start;
		MetadataReader reader;
		Section code_section;

		MethodDefinition method;
		MethodBody body;
		Instruction [] instruction_map;

		int Offset {
			get { return base.position - start; }
		}

		CodeReader (Section section, MetadataReader reader)
			: base (section.Data)
		{
			this.code_section = section;
			this.reader = reader;
		}

		public static CodeReader CreateCodeReader (MetadataReader metadata)
		{
			return new CodeReader (metadata.image.MetadataSection, metadata);
		}

		public MethodBody ReadMethodBody (MethodDefinition method)
		{
			this.method = method;
			this.body = new MethodBody (method);

			reader.context = method;

			ReadMethodBody ();

			return this.body;
		}

		void MoveTo (int rva)
		{
			if (!IsInSection (rva)) {
				code_section = reader.image.GetSectionAtVirtualAddress ((uint) rva);
				Reset (code_section.Data);
			}

			base.position = rva - (int) code_section.VirtualAddress;
		}

		bool IsInSection (int rva)
		{
			return code_section.VirtualAddress <= rva && rva < code_section.VirtualAddress + code_section.SizeOfRawData;
		}

		void ReadMethodBody ()
		{
			MoveTo (method.RVA);

			var flags = ReadByte ();
			switch (flags & 0x3) {
			case 0x2: // tiny
				body.CodeSize = flags >> 2;
				body.MaxStackSize = 8;
				ReadCode ();
				break;
			case 0x3: // fat
				base.position--;
				ReadFatMethod ();
				break;
			default:
				throw new InvalidOperationException ();
			}

			var symbol_reader = reader.module.SymbolReader;

			if (symbol_reader != null)
				symbol_reader.Read (body, instruction_map);
		}

		void ReadFatMethod ()
		{
			var flags = ReadUInt16 ();
			body.MaxStackSize = ReadUInt16 ();
			body.CodeSize = (int) ReadUInt32 ();
			body.LocalVarToken = new MetadataToken (ReadUInt32 ());
			body.InitLocals = (flags & 0x10) != 0;

			if (body.LocalVarToken.RID != 0)
				ReadVariables ();

			ReadCode ();

			if ((flags & 0x8) != 0)
				ReadSection ();
		}

		void ReadVariables ()
		{
			var position = reader.position;
			reader.ReadVariables (body);
			reader.position = position;
		}

		void AdjustInstructionMap (int size)
		{
			if (instruction_map == null || size > instruction_map.Length)
				instruction_map = new Instruction [size];
			else
				Array.Clear (instruction_map, 0, instruction_map.Length);
		}

		void ReadCode ()
		{
			start = position;
			var end = start + body.CodeSize;
			Instruction previous = null;
			var instructions = body.Instructions;

			AdjustInstructionMap (body.CodeSize);

			while (position < end) {
				var offset = Offset;
				var opcode = ReadOpCode ();
				var current = new Instruction (offset, opcode);
				instruction_map [current.Offset] = current;
				if (previous != null) {
					current.Previous = previous;
					previous.Next = current;
				}

				if (opcode.OperandType != OperandType.InlineNone)
					current.Operand = ReadOperand (current);

				instructions.Add (current);

				previous = current;
			}

			ResolveBranches (instructions);
		}

		OpCode ReadOpCode ()
		{
			var il_opcode = ReadByte ();
			return il_opcode != 0xfe
			  ? OpCodes.OneByteOpCode [il_opcode]
			  : OpCodes.TwoBytesOpCode [ReadByte ()];
		}

		object ReadOperand (Instruction instruction)
		{
			switch (instruction.OpCode.OperandType) {
			case OperandType.InlineSwitch:
				var length = ReadInt32 ();
				var base_offset = Offset + (4 * length);
				var branches = new int [length];
				for (int i = 0; i < length; i++)
					branches [i] = base_offset + ReadInt32 ();
				return branches;
			case OperandType.ShortInlineBrTarget:
				return ReadSByte () + Offset;
			case OperandType.InlineBrTarget:
				return ReadInt32 () + Offset;
			case OperandType.ShortInlineI:
				if (instruction.OpCode == OpCodes.Ldc_I4_S)
					return (sbyte) ReadByte ();
				else
					return ReadByte ();
			case OperandType.InlineI:
				return ReadInt32 ();
			case OperandType.ShortInlineR:
				return ReadSingle ();
			case OperandType.InlineR:
				return ReadDouble ();
			case OperandType.InlineI8:
				return ReadInt64 ();
			case OperandType.ShortInlineVar:
				return GetVariable (ReadByte ());
			case OperandType.InlineVar:
				return GetVariable (ReadUInt16 ());
			case OperandType.ShortInlineParam:
				return GetParameter (ReadByte ());
			case OperandType.InlineParam:
				return GetParameter (ReadUInt16 ());
			case OperandType.InlineSig:
				return GetCallSite (ReadToken ());
			case OperandType.InlineString:
				return reader.image.UserStringHeap [ReadToken ().RID];
			case OperandType.InlineTok:
			case OperandType.InlineType:
			case OperandType.InlineMethod:
			case OperandType.InlineField:
				return reader.LookupToken (ReadToken ());
			default:
				throw new NotSupportedException ();
			}
		}

		ParameterDefinition GetParameter (int index)
		{
			if (method.HasThis) {
				if (index == 0)
					return body.ThisParameter;
				index--;
			}

			return method.Parameters [index];
		}

		VariableDefinition GetVariable (int index)
		{
			return body.Variables [index];
		}

		CallSite GetCallSite (MetadataToken token)
		{
			return reader.ReadCallSite (token);
		}

		void ResolveBranches (InstructionCollection instructions)
		{
			for (int i = 0; i < instructions.Count; i++) {
				var instruction = instructions [i];
				switch (instruction.OpCode.OperandType) {
				case OperandType.ShortInlineBrTarget:
				case OperandType.InlineBrTarget:
					instruction.Operand = GetInstruction ((int) instruction.Operand);
					break;
				case OperandType.InlineSwitch:
					var offsets = (int []) instruction.Operand;
					var branches = new Instruction [offsets.Length];
					for (int j = 0; j < offsets.Length; j++)
						branches [j] = GetInstruction (offsets [j]);

					instruction.Operand = branches;
					break;
				}
			}
		}

		Instruction GetInstruction (int offset)
		{
			if (offset < 0 || offset > instruction_map.Length)
				return null;

			return instruction_map [offset];
		}

		void ReadSection ()
		{
			Align (4);

			const byte fat_format = 0x40;
			const byte more_sects = 0x80;

			var flags = ReadByte ();
			if ((flags & fat_format) == 0)
				ReadSmallSection ();
			else
				ReadFatSection ();

			if ((flags & more_sects) != 0)
				ReadSection ();
		}

		void ReadSmallSection ()
		{
			int length = ReadByte () / 12;
			Advance (2);

			ReadExceptionHandlers (
				length,
				() => (int) ReadUInt16 (),
				() => (int) ReadByte ());
		}

		void ReadFatSection ()
		{
			position--;
			int length = (ReadInt32 () >> 8) / 24;

			ReadExceptionHandlers (
				length,
				ReadInt32,
				ReadInt32);
		}

		// inline ?
		void ReadExceptionHandlers (int count, Func<int> read_entry, Func<int> read_length)
		{
			for (int i = 0; i < count; i++) {
				var handler = new ExceptionHandler (
					(ExceptionHandlerType) (read_entry () & 0x7));

				handler.TryStart = GetInstruction (read_entry ());
				handler.TryEnd = GetInstruction (handler.TryStart.Offset + read_length ());

				handler.HandlerStart = GetInstruction (read_entry ());
				handler.HandlerEnd = GetInstruction (handler.HandlerStart.Offset + read_length ());

				ReadExceptionHandlerSpecific (handler);

				this.body.ExceptionHandlers.Add (handler);
			}
		}

		void ReadExceptionHandlerSpecific (ExceptionHandler handler)
		{
			switch (handler.HandlerType) {
			case ExceptionHandlerType.Catch:
				handler.CatchType = (TypeReference) reader.LookupToken (ReadToken ());
				break;
			case ExceptionHandlerType.Filter:
				handler.FilterStart = GetInstruction (ReadInt32 ());
				handler.FilterEnd = handler.HandlerStart.Previous;
				break;
			default:
				Advance (4);
				break;
			}
		}

		void Align (int align)
		{
			align--;
			Advance (((position + align) & ~align) - position);
		}

		MetadataToken ReadToken ()
		{
			return new MetadataToken (ReadUInt32 ());
		}
	}
}
